/*
 * @Author: wuqingchun
 * @Date: 2024-03-28 17:05:48
 * @Description: 
 * @LastEditors: wuqingchun
 * @LastEditTime: 2024-04-25 09:25:13
 */
#ifndef __PERF_LOCK_H__
#define __PERF_LOCK_H__

#include <vector>
#include <atomic>
#include <thread>
#include <mutex>

#include <unistd.h>


class AtomicMutex
{
private:
    /* data */

    std::atomic_flag m_objFlag = ATOMIC_FLAG_INIT;
public:
    void lock()
    {
        while (!m_objFlag.test_and_set(std::memory_order_acquire))
        {
            
        }
    }

    void unlock()
    {
        m_objFlag.clear(std::memory_order_release);
    }
};


class CLangMutex
{
private:
    pthread_mutex_t m_stMtx;
public:
    CLangMutex()
    {
        (void)pthread_mutex_init(&m_stMtx, nullptr);
    }

    void lock()
    {
        (void)pthread_mutex_lock(&m_stMtx);
    }

    void unlock()
    {
        (void)pthread_mutex_unlock(&m_stMtx);
    }
};

#if defined USE_ATOMIC
using MUTEX = AtomicMutex;
#elif  defined USE_CLANG
using MUTEX = CLangMutex;
#else
using MUTEX = std::mutex;
#endif


// 锁性能测试APP
class LockPerfTestApp
{
    struct ThreadInfo
    {
        std::thread Thread;
        uint64_t LockUsedTimeNS;
        uint64_t UnlockUsedTimeNS;
        uint64_t ExecuteCount;
    };
private:
    uint64_t m_uiCount {0};
    volatile bool m_bRunningFlag {false};
    std::vector<ThreadInfo> m_vecThreads;
    MUTEX m_objMtx;

    AtomicMutex m_objAtomicMutex;
    std::mutex m_objMutex;
    CLangMutex m_objCMutex;
    std::string m_strLockType;
public:
    void Start(uint32_t uiThreadNum, std::string& refLockType)
    {
        m_strLockType = refLockType;

        m_bRunningFlag = true;
        for (uint32_t uiTmp = 0; uiTmp < uiThreadNum; uiTmp++)
        {
            //m_vecThreads.push_back({std::thread(&LockPerfTestApp::__ThreadLockTestFunc, this, uiTmp), 0, 0, 0});

            if (refLockType == "atomic lock")
            {
                m_vecThreads.push_back({std::thread(&LockPerfTestApp::__ThreadAtomicFunc, this, uiTmp), 0, 0, 0});
            }
            else if (refLockType == "cpp mutex")
            {
                m_vecThreads.push_back({std::thread(&LockPerfTestApp::__ThreadCPPFunc, this, uiTmp), 0, 0, 0});
            }
            else
            {
                m_vecThreads.push_back({std::thread(&LockPerfTestApp::__ThreadCLangFunc, this, uiTmp), 0, 0, 0});
            }
        }
    }

    void Stop()
    {
        m_bRunningFlag = false;

        for (auto& refItem : m_vecThreads)
        {
            if (refItem.Thread.joinable())
            {
                refItem.Thread.join();
            }
        }

        // 统计耗时

        uint64_t ui64TotalLockCostTime = 0;
        uint64_t ui64TotalUnlockCostTime = 0;
        uint64_t ui64TotalExecuteCount = 0;
        for (auto& refItem : m_vecThreads)
        {
            ui64TotalLockCostTime += refItem.LockUsedTimeNS;
            ui64TotalUnlockCostTime += refItem.UnlockUsedTimeNS;
            ui64TotalExecuteCount += refItem.ExecuteCount;
        }

        printf("LockType: %s\n", m_strLockType.c_str());
        printf("LockCostTime: %lu\n", ui64TotalLockCostTime/ui64TotalExecuteCount);
        printf("UnlockCostTime: %lu\n", ui64TotalUnlockCostTime/ui64TotalExecuteCount);

        m_vecThreads.clear();
    }


private:
    void __ThreadLockTestFunc(uint32_t uiThreadID)
    {
        __SetAffinity(uiThreadID + 1);

        while (m_bRunningFlag)
        {
            auto tsStart = std::chrono::steady_clock::now().time_since_epoch().count();
            m_objMtx.lock();
            auto tsEnd = std::chrono::steady_clock::now().time_since_epoch().count();

            m_vecThreads[uiThreadID].LockUsedTimeNS += tsEnd - tsStart;

            m_vecThreads[uiThreadID].ExecuteCount++;
    
            tsStart = std::chrono::steady_clock::now().time_since_epoch().count();
            m_objMtx.unlock();
            tsEnd = std::chrono::steady_clock::now().time_since_epoch().count();

            m_vecThreads[uiThreadID].UnlockUsedTimeNS += tsEnd - tsStart;
        }
    }

    void __ThreadAtomicFunc(uint32_t uiThreadID)
    {
        __SetAffinity(uiThreadID + 1);

        while (m_bRunningFlag)
        {
            auto tsStart = std::chrono::steady_clock::now().time_since_epoch().count();
            m_objAtomicMutex.lock();
            auto tsEnd = std::chrono::steady_clock::now().time_since_epoch().count();

            m_vecThreads[uiThreadID].LockUsedTimeNS += tsEnd - tsStart;

            m_vecThreads[uiThreadID].ExecuteCount++;
    
            tsStart = std::chrono::steady_clock::now().time_since_epoch().count();
            m_objAtomicMutex.unlock();
            tsEnd = std::chrono::steady_clock::now().time_since_epoch().count();

            m_vecThreads[uiThreadID].UnlockUsedTimeNS += tsEnd - tsStart;
        }
    }

    void __ThreadCLangFunc(uint32_t uiThreadID)
    {
        __SetAffinity(uiThreadID + 1);

        while (m_bRunningFlag)
        {
            auto tsStart = std::chrono::steady_clock::now().time_since_epoch().count();
            m_objCMutex.lock();
            auto tsEnd = std::chrono::steady_clock::now().time_since_epoch().count();

            m_vecThreads[uiThreadID].LockUsedTimeNS += tsEnd - tsStart;

            m_vecThreads[uiThreadID].ExecuteCount++;
    
            tsStart = std::chrono::steady_clock::now().time_since_epoch().count();
            m_objCMutex.unlock();
            tsEnd = std::chrono::steady_clock::now().time_since_epoch().count();

            m_vecThreads[uiThreadID].UnlockUsedTimeNS += tsEnd - tsStart;
        }
    }

    void __ThreadCPPFunc(uint32_t uiThreadID)
    {
        __SetAffinity(uiThreadID + 1);

        while (m_bRunningFlag)
        {
            auto tsStart = std::chrono::steady_clock::now().time_since_epoch().count();
            m_objMutex.lock();
            auto tsEnd = std::chrono::steady_clock::now().time_since_epoch().count();

            m_vecThreads[uiThreadID].LockUsedTimeNS += tsEnd - tsStart;

            m_vecThreads[uiThreadID].ExecuteCount++;
    
            tsStart = std::chrono::steady_clock::now().time_since_epoch().count();
            m_objMutex.unlock();
            tsEnd = std::chrono::steady_clock::now().time_since_epoch().count();

            m_vecThreads[uiThreadID].UnlockUsedTimeNS += tsEnd - tsStart;
        }
    }

    void __SetAffinity(int32_t iCpu)
    {
        cpu_set_t cpuset;
        CPU_ZERO(&cpuset);
        CPU_SET(iCpu, &cpuset);

        (void)pthread_setaffinity_np(pthread_self(), sizeof(cpu_set_t), &cpuset);
    }
};

#endif //__PERF_LOCK_H__